#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
敏感词检测
"""
import re

from apps.health.models import SensitiveWord


class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False
        self.category = None

class SensitiveWordTrieDetector:
    def __init__(self):
        self.root = TrieNode()
        self._is_initialized = False

    def initialize(self):
        """从数据库加载敏感词构建Trie树"""
        if self._is_initialized:
            return

        sensitive_words = SensitiveWord.objects.all()

        for word_obj in sensitive_words:
            self._insert_word(word_obj.word, word_obj.category)

        self._is_initialized = True

    def _insert_word(self, word, category):
        """将敏感词插入Trie树"""
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.is_end = True
        node.category = category

    def detect(self, text):
        """检测文本中的敏感词"""
        if not self._is_initialized:
            raise Exception('SensitiveWordTrieDetector not initialized. Call initialize() first.')

        sensitive_words = []
        text_length = len(text)

        for i in range(text_length):
            node = self.root
            j = i
            while j < text_length and text[j] in node.children:
                node = node.children[text[j]]
                if node.is_end:
                    sensitive_words.append({
                        'word': text[i:j+1],
                        'category': node.category
                    })
                j += 1

        return len(sensitive_words) > 0, [word['word'] for word in sensitive_words]

    def filter(self, text, replacement="***"):
        """过滤文本中的敏感词"""
        is_sensitive, sensitive_words = self.detect(text)
        if not is_sensitive:
            return text

        filtered_text = text
        for word in sensitive_words:
            escaped_word = re.escape(word)
            filtered_text = re.sub(escaped_word, replacement, filtered_text)

        return filtered_text
